// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/dwio/batch_maker.h>
#include <pollux/exec/plan_node_stats.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/parse/expressions.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;

using kumo::pollux::test::BatchMaker;

class FilterProjectTest : public OperatorTestBase {
 protected:
  void assertFilter(
      const std::vector<RowVectorPtr>& vectors,
      const std::string& filter = "c1 % 10  > 0") {
    auto plan = PlanBuilder().values(vectors).filter(filter).planNode();

    assertQuery(plan, "SELECT * FROM tmp WHERE " + filter);
  }

  void assertProject(const std::vector<RowVectorPtr>& vectors) {
    auto plan = PlanBuilder()
                    .values(vectors)
                    .project({"c0", "c1", "c0 + c1"})
                    .planNode();

    auto task = assertQuery(plan, "SELECT c0, c1, c0 + c1 FROM tmp");

    // A quick sanity check for memory usage reporting. Check that peak total
    // memory usage for the project node is > 0.
    auto planStats = toPlanStats(task->taskStats());
    auto projectNodeId = plan->id();
    auto it = planStats.find(projectNodeId);
    ASSERT_TRUE(it != planStats.end());
    ASSERT_TRUE(it->second.peakMemoryBytes > 0);
  }

  std::shared_ptr<const RowType> rowType_{
      ROW({"c0", "c1", "c2", "c3"},
          {BIGINT(), INTEGER(), SMALLINT(), DOUBLE()})};
};

TEST_F(FilterProjectTest, filter) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));
    vectors.push_back(vector);
  }
  createDuckDbTable(vectors);

  assertFilter(vectors);
}

TEST_F(FilterProjectTest, filterOverDictionary) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));

    auto indices =
        AlignedBuffer::allocate<int32_t>(2 * vector->size(), pool_.get());
    auto indicesPtr = indices->asMutable<int32_t>();
    for (int32_t j = 0; j < vector->size() / 2; j++) {
      indicesPtr[2 * j] = j;
      indicesPtr[2 * j + 1] = j;
    }
    std::vector<VectorPtr> newChildren = vector->children();
    newChildren[1] = BaseVector::wrap_in_dictionary(
        BufferPtr(nullptr), indices, vector->size(), vector->childAt(1));
    vectors.push_back(std::make_shared<RowVector>(
        pool_.get(),
        rowType_,
        BufferPtr(nullptr),
        vector->size(),
        newChildren,
        0 /*nullCount*/));
  }
  createDuckDbTable(vectors);

  assertFilter(vectors);
}

TEST_F(FilterProjectTest, filterOverConstant) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));

    std::vector<VectorPtr> newChildren = vector->children();
    newChildren[1] =
        BaseVector::wrap_in_constant(vector->size(), 7, vector->childAt(1));
    vectors.push_back(std::make_shared<RowVector>(
        pool_.get(),
        rowType_,
        BufferPtr(nullptr),
        vector->size(),
        newChildren,
        0 /*nullCount*/));
  }
  createDuckDbTable(vectors);

  assertFilter(vectors);
}

TEST_F(FilterProjectTest, project) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));
    vectors.push_back(vector);
  }
  createDuckDbTable(vectors);

  assertProject(vectors);
}

TEST_F(FilterProjectTest, projectOverDictionary) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));

    auto indices =
        AlignedBuffer::allocate<int32_t>(2 * vector->size(), pool_.get());
    auto indicesPtr = indices->asMutable<int32_t>();
    for (int32_t j = 0; j < vector->size() / 2; j++) {
      indicesPtr[2 * j] = j;
      indicesPtr[2 * j + 1] = j;
    }
    std::vector<VectorPtr> newChildren = vector->children();
    newChildren[1] = BaseVector::wrap_in_dictionary(
        BufferPtr(nullptr), indices, vector->size(), vector->childAt(1));
    vectors.push_back(std::make_shared<RowVector>(
        pool_.get(),
        rowType_,
        BufferPtr(nullptr),
        vector->size(),
        newChildren,
        0 /*nullCount*/));
  }
  createDuckDbTable(vectors);

  assertProject(vectors);
}

TEST_F(FilterProjectTest, projectOverConstant) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));

    std::vector<VectorPtr> newChildren = vector->children();
    newChildren[1] =
        BaseVector::wrap_in_constant(vector->size(), 7, vector->childAt(1));
    vectors.push_back(std::make_shared<RowVector>(
        pool_.get(),
        rowType_,
        BufferPtr(nullptr),
        vector->size(),
        newChildren,
        0 /*nullCount*/));
  }
  createDuckDbTable(vectors);

  assertProject(vectors);
}

TEST_F(FilterProjectTest, projectOverLazy) {
  vector_size_t size = 1'000;
  auto valueAtC0 = [](auto row) -> int32_t {
    return row % 2 == 0 ? row : -row;
  };
  auto valueAtC1 = [](auto row) -> double {
    return row % 3 == 0 ? row * 0.1 : -row * 0.1;
  };
  auto lazyVectors = make_row_vector({
      vectorMaker_.lazy_flat_vector<int32_t>(size, valueAtC0),
      vectorMaker_.lazy_flat_vector<double>(size, valueAtC1),
  });

  auto vectors = make_row_vector({
      make_flat_vector<int32_t>(size, valueAtC0),
      make_flat_vector<double>(size, valueAtC1),
  });

  createDuckDbTable({vectors});

  auto plan = PlanBuilder()
                  .values({lazyVectors})
                  .project({"c0 > 0 AND c1 > 0.0", "c1 + 5.2"})
                  .planNode();
  assertQuery(plan, "SELECT c0 > 0 AND c1 > 0, c1 + 5.2 FROM tmp");
}

TEST_F(FilterProjectTest, filterProject) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));
    vectors.push_back(vector);
  }
  createDuckDbTable(vectors);

  auto plan = PlanBuilder()
                  .values(vectors)
                  .filter("c1 % 10  > 0")
                  .project({"c0", "c1", "c0 + c1"})
                  .planNode();

  assertQuery(plan, "SELECT c0, c1, c0 + c1 FROM tmp WHERE c1 % 10 > 0");
}

TEST_F(FilterProjectTest, dereference) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));
    vectors.push_back(vector);
  }
  createDuckDbTable(vectors);

  auto plan = PlanBuilder()
                  .values(vectors)
                  .project({"row_constructor(c1, c2) AS c1_c2"})
                  .project({"c1_c2.c1", "c1_c2.c2"})
                  .planNode();
  assertQuery(plan, "SELECT c1, c2 FROM tmp");

  plan = PlanBuilder()
             .values(vectors)
             .project({"row_constructor(c1, c2) AS c1_c2"})
             .filter("c1_c2.c1 % 10 = 5")
             .project({"c1_c2.c1", "c1_c2.c2"})
             .planNode();
  assertQuery(plan, "SELECT c1, c2 FROM tmp WHERE c1 % 10 = 5");
}

TEST_F(FilterProjectTest, allFailedOrPassed) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    // We alternate between a batch where all pass and a batch where
    // no row passes. c0 is flat vector. c1 is constant vector.
    const int32_t value = i % 2 == 0 ? 0 : 1;

    vectors.push_back(make_row_vector({
        make_flat_vector<int32_t>(100, [&](auto row) { return value; }),
        make_constant(value, 100),
    }));
  }
  createDuckDbTable(vectors);

  // filter over flat vector
  assertFilter(vectors, "c0 = 0");

  // filter over constant vector
  assertFilter(vectors, "c1 = 0");
}

// Tests fusing of consecutive filters and projects.
TEST_F(FilterProjectTest, filterProjectFused) {
  std::vector<RowVectorPtr> vectors;
  for (int32_t i = 0; i < 10; ++i) {
    auto vector = std::dynamic_pointer_cast<RowVector>(
        BatchMaker::createBatch(rowType_, 100, *pool_));
    vectors.push_back(vector);
  }
  createDuckDbTable(vectors);

  auto plan = PlanBuilder()
                  .values(vectors)
                  .filter("c0 % 10 < 9")
                  .project({"c0", "c1", "c0 % 100 + c1 % 50 AS e1"})
                  .filter("c0 % 10 < 8")
                  .project({"c0", "c1", "e1", "c0 % 100 AS e2"})
                  .filter("c0 % 10 < 5")
                  .project({"c0", "c1", "e1", "e2"})
                  .planNode();

  assertQuery(
      plan,
      "SELECT c0, c1, c0 %100 + c1 % 50, c0 % 100 FROM tmp WHERE c0 % 10 < 5");
}

TEST_F(FilterProjectTest, projectAndIdentityOverLazy) {
  // Verify that a lazy column which is a part of both an identity projection
  // and a regular projection is loaded correctly. This is done by running a
  // projection that only operates over a subset of the rows of the lazy vector.
  vector_size_t size = 20;
  auto value_at = [](auto row) -> int32_t { return row; };
  auto lazyVectors = make_row_vector({
      make_flat_vector<int32_t>(size, value_at),
      vectorMaker_.lazy_flat_vector<int32_t>(size, value_at),
  });

  auto vectors = make_row_vector({
      make_flat_vector<int32_t>(size, value_at),
      make_flat_vector<int32_t>(size, value_at),
  });

  createDuckDbTable({vectors});

  auto plan = PlanBuilder()
                  .values({lazyVectors})
                  .project({"c0 < 10 AND c1 < 10", "c1"})
                  .planNode();
  assertQuery(plan, "SELECT c0 < 10 AND c1 < 10, c1 FROM tmp");
}

// Verify the optimization of avoiding copy in null propagation does not break
// the case when the field is shared between multiple parents.
TEST_F(FilterProjectTest, nestedFieldReferenceSharedChild) {
  auto shared = make_flat_vector<int64_t>(10, melon::identity);
  auto vector = make_row_vector({
      make_row_vector({
          make_row_vector({shared}, null_every(2)),
          make_row_vector({shared}, null_every(3)),
      }),
  });
  auto plan =
      PlanBuilder()
          .values({vector})
          .project({"coalesce((c0).c0.c0, 0) + coalesce((c0).c1.c0, 0)"})
          .planNode();
  auto expected = make_flat_vector<int64_t>(10);
  for (int i = 0; i < 10; ++i) {
    expected->set(i, (i % 2 == 0 ? 0 : i) + (i % 3 == 0 ? 0 : i));
  }
  AssertQueryBuilder(plan).assertResults(make_row_vector({expected}));
}

TEST_F(FilterProjectTest, numSilentThrow) {
  auto row = make_row_vector(
      {make_flat_vector<int32_t>(100, [&](auto row) { return row; })});

  core::PlanNodeId filterId;
  // Change the plan when /0 error is fixed not to throw.
  auto plan = PlanBuilder()
                  .values({row})
                  .filter("try (c0 / 0) = 1")
                  .capturePlanNodeId(filterId)
                  .planNode();

  auto task = AssertQueryBuilder(plan).assertEmptyResults();
  auto planStats = toPlanStats(task->taskStats());
  ASSERT_EQ(100, planStats.at(filterId).customStats.at("numSilentThrow").sum);
}

TEST_F(FilterProjectTest, statsSplitter) {
  auto data = make_row_vector({
      make_flat_vector<int32_t>(100, melon::identity),
  });

  // Make 5 batches, 20 rows each: 0...19, 20...39, 40...59, 60...79, 80...99.
  // Filter out firs 25 rows (1 full batch plus some more), then keep every 3rd
  // row.
  core::PlanNodeId filterId;
  core::PlanNodeId projectId;
  auto plan = PlanBuilder()
                  .values(split(data, 5))
                  .filter("if (c0 < 25, false, c0 % 3 = 0)")
                  .capturePlanNodeId(filterId)
                  .project({"c0", "c0 + 1"})
                  .capturePlanNodeId(projectId)
                  .planNode();

  std::shared_ptr<Task> task;
  AssertQueryBuilder(plan).runWithoutResults(task);

  auto planStats = toPlanStats(task->taskStats());

  const auto& filterStats = planStats.at(filterId);
  const auto& projectStats = planStats.at(projectId);

  EXPECT_EQ(filterStats.inputRows, 100);
  EXPECT_EQ(filterStats.inputVectors, 5);

  EXPECT_EQ(filterStats.outputRows, 25);
  EXPECT_EQ(filterStats.outputVectors, 4);
  EXPECT_LT(filterStats.outputBytes * 2, filterStats.inputBytes);

  EXPECT_EQ(projectStats.inputRows, 25);
  EXPECT_EQ(projectStats.inputVectors, 4);

  EXPECT_EQ(projectStats.inputBytes, filterStats.outputBytes);
  EXPECT_LT(projectStats.inputBytes, projectStats.outputBytes);

  EXPECT_EQ(projectStats.outputRows, 25);
  EXPECT_EQ(projectStats.outputVectors, 4);
}
